import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# --- 唯一需要修改的地方 ---
# 将模型名称替换为你的本地文件夹路径
# Windows 示例: "D:/models/finbert"
# Linux/macOS 示例: "/home/user/models/finbert"
local_model_path = "/Users/zhangyalei/work_space/model_workspace/finbert"
# -------------------------

print(f"正在从本地路径 {local_model_path} 加载模型和分词器...")

# 加载分词器，现在它会从本地路径读取
tokenizer = AutoTokenizer.from_pretrained(local_model_path)

# 加载模型，现在它会从本地路径读取权重
model = AutoModelForSequenceClassification.from_pretrained(local_model_path)

print("模型加载完成！")

# --- 推理部分与之前完全相同 ---

# 定义一些待分析的金融新闻标题或句子
texts = [
    "The company's quarterly earnings exceeded all expectations, sending the stock soaring.",
    "Investors are concerned about the potential impact of new regulations on the tech sector.",
    "The market remained stable as traders awaited the latest inflation report."
]

# 获取模型的标签映射
id2label = model.config.id2label
print(f"\n模型标签映射: {id2label}\n")

# 对每个文本进行情感分析
for text in texts:
    print(f"分析文本: '{text}'")
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)

    with torch.no_grad():
        logits = model(**inputs).logits

    predictions = torch.argmax(logits, dim=-1)
    predicted_class_id = predictions.item()
    predicted_label = id2label[predicted_class_id]

    probabilities = torch.softmax(logits, dim=-1).squeeze().tolist()
    scores = {id2label[i]: prob for i, prob in enumerate(probabilities)}

    print(f"  -> 预测情感: {predicted_label}")
    print(f"  -> 详细分数: {scores}\n")

